################################################################################
#
# NIST/NSRL Loader
#
################################################################################

import pyflag.DB as DB
import pyflag.FlagFramework as FlagFramework
import pyflag.Reports as Reports
import pyflag.conf
import pyflag.pyflaglog as pyflaglog
import re
from SoftwareUtil import init_index

config = pyflag.conf.ConfObject()
config.add_option("NSRL_DB", default = "nsrldb")
config.add_option("NSRL_HASHES", default = "hashes")
config.add_option("NSRL_PRODUCTS", default = "products")
config.add_option("NSRL_OS", default = "os")
config.add_option("NSRL_MFG", default = "mfg")



# Process NSRLProd.txt file
def processProductLine(dbh, line):
    result = re.compile(r"(?P<ProductCode>\d+),\"(?P<ProductName>.*)\",\"(?P<ProductVersion>.*)\",\"(?P<OpSystemCode>.*)\",\"(?P<MfgCode>.*)\",\"(?P<Language>.*)\",\"(?P<AppType>.*)\"")
    matching = result.match(line.strip())

    if matching != None:
	dbh.insert(config.NSRL_PRODUCTS,
	    ProductCode = matching.group('ProductCode'),
	    ProductName = matching.group('ProductName'),
	    ProductVersion = matching.group('ProductVersion'),
	    OpSystemCode = matching.group('OpSystemCode'),
	    MfgCode = matching.group('MfgCode'),
	    Language = matching.group('Language'),
	    ApplicationType = matching.group('AppType'))
    else:
	pyflaglog.log(pyflaglog.ERROR, "Invalid format line: %s" % (line))



# Process NSRLFile.txt file
def processFilesLine(dbh, line):
    result = re.compile(r"\"(?P<SHA1>.*)\",\"(?P<MD5>.*)\",\"(?P<CRC32>.*)\",\"(?P<FileName>.*)\",(?P<FileSize>\d+),(?P<ProductCode>\d+),\"(?P<OpSystemCode>.*)\",\"(?P<SpecialCode>.*)\"")
    matching = result.match(line.strip())

    if matching != None:
        dbh.insert(config.NSRL_HASHES,
	    SHA1 = matching.group('SHA1'),
	    MD5 = matching.group('MD5'),
	    CRC32 = matching.group('CRC32'),
	    FileName = matching.group('FileName'),
	    FileSize = matching.group('FileSize'),
	    ProductCode = matching.group('ProductCode'),
	    OpSystemCode = matching.group('OpSystemCode'),
	    SpecialCode = matching.group('SpecialCode'))
    else:
	pyflaglog.log(pyflaglog.ERROR, "Invalid format line: %s" % (line))



class NSRLTables(FlagFramework.EventHandler):
    def startup(self):
	try:
	    dbh = DB.DBO("information_schema")
	    dbh.execute("""SELECT * FROM TABLES WHERE TABLE_SCHEMA = '%s'""", config.NSRL_DB)
	    row = dbh.fetch()

	    if row == None: self.init_default_db(None, None)
	except:
	    pass


    def init_default_db(self, dbh, case):
	try:
	    dbh = DB.DBO()
	    dbh.execute("""CREATE DATABASE IF NOT EXISTS %s""" % config.NSRL_DB)
	    dbh.execute("""CREATE TABLE IF NOT EXISTS `%s`.`%s` (
        	`SID` int(16) NOT NULL auto_increment,
        	`SHA1` char(40) NOT NULL,
        	`MD5` char(32) NOT NULL,
		`CRC32` char(8) NOT NULL,
		`FileName` varchar(255) NOT NULL,
		`FileSize` int(16) NOT NULL,
        	`ProductCode` int(16) NOT NULL default '0',
		`OpSystemCode` varchar(15) NOT NULL,
		`SpecialCode` char(1),
        	PRIMARY KEY (`SID`),
		KEY sha1_index (`SHA1`))""" % (config.NSRL_DB, config.NSRL_HASHES))
	    dbh.execute("""CREATE TABLE IF NOT EXISTS `%s`.`%s` (
		`PID` int(16) NOT NULL auto_increment,
        	`ProductCode` int(16) NOT NULL,
        	`ProductName` varchar(150) NOT NULL,
        	`ProductVersion` varchar(15) NOT NULL,
        	`OpSystemCode` varchar(15) NOT NULL default '0',
        	`MfgCode` varchar(15) NOT NULL default '0',
        	`Language`varchar(150) NOT NULL,
        	`ApplicationType` varchar(50) NOT NULL,
        	PRIMARY KEY (`PID`))""" % (config.NSRL_DB, config.NSRL_PRODUCTS))
	    ## TO DO MFG & OS TABLES

	    dbh.execute("""CREATE TABLE IF NOT EXISTS `%s`.`%s` (
		`SHA1` char(40) NOT NULL,
		`MD5` char(32) NOT NULL,
		`Count` INT(16) NOT NULL,
		KEY sha1_i (`SHA1`))""" % (config.NSRL_DB, config.NSRL_HASH_COUNT))
	    
	    dbh.execute("""CREATE TABLE IF NOT EXISTS `%s`.`%s` (
		`ProductCode` int(16) NOT NULL,
		`NSRLCount` int(16) NOT NULL,
		`NSRLMin` double NOT NULL,
		`NSRLNormal` double NOT NULL,
		`NSRLExtra` double NOT NULL,
		`NSRLUltra` double NOT NULL,
		`NSRL_LIN_15` double NOT NULL,
		`NSRL_LIN_10` double NOT NULL,
		`NSRL_LIN_5` double NOT NULL,
		KEY pc_i (`ProductCode`))""" % (config.NSRL_DB, config.NSRL_PRODUCT_COUNT))
	except:
	    pass


class NSRLLoad(Reports.report):
    parameters = {"A":"any", "B":"any", "C":"any", "D":"any", "prod":"any", "final":"string"}
    name = "NIST/NSRL data load"
    family = "Hashes"
    order = 10


    def __init__(self, flag = None, ui = None):
	Reports.report.__init__(self, flag, ui)
	self.dbh = DB.DBO(config.NSRL_DB)
	self.__progress = True


    def form(self, query, result):
	result.fileselector("Select NSRLFile from A disk", "A")
	result.fileselector("Select NSRLFile from B disk", "B")
	result.fileselector("Select NSRLFile from C disk", "C")
	result.fileselector("Select NSRLFile from D disk", "D")
	result.fileselector("Select NSRLProduct", "prod")
	result.checkbox("Click here when finished", "final", "ok")


    def progress(self, query, result):
	if self.__progress:
	    self.dbh.execute("""SELECT count(*) AS count FROM %s""" % config.NSRL_HASHES)
	    tmp = self.dbh.fetch()
	    result.heading("Loading...")
	    #result.para("Currently load %d lines" % tmp['count'])
	else:
	    result.heading("Loading...")
	    result.para("Creating useful tables")


    def display(self, query, result):
	#self.dbh.expire_cache()
	result.heading("NSRLFiles file loaded successful")


    def analyse(self, query):
	self.dbh.execute("""DELETE FROM %s""" % config.NSRL_HASHES)
	self.dbh.execute("""DELETE FROM %s""" % config.NSRL_PRODUCTS)
	self.dbh.execute("""DELETE FROM %s""" % config.NSRL_HASH_COUNT)
	self.dbh.execute("""DELETE FROM %s""" % config.NSRL_PRODUCT_COUNT)
	filelist = [query['A'], query['B'], query['C'], query['D']]

	## Read NSRLFiles.txt
	for filename in filelist:
	    try:
		print filename
		file = None
		file = open(filename, 'r')
		line = file.readline()
		line = file.readline()

		while line:
		    processFilesLine(self.dbh, line)
		    line = file.readline()
	    except IOError, e:
		pyflaglog.log(pyflaglog.ERROR, str(e))
	    finally:
		if file != None:
		    file.close()
	
	## Read product files
	try:
	    filename = query['prod']
	    file = open(filename)
	    lines = file.readlines()
	    print len(lines)
	    for line in lines[1:]:
		processProductLine(self.dbh, line)
	except IOError, e:
	    pyflaglog.log(pyflaglog.ERROR, str(e))
	finally:
	    if file != None:
		file.close()

	init_index(self.dbh)
